#include "MyNpzScoreWriter.hh"

#include <string>
#include <vector>

#include "G4StatDouble.hh"
#include "G4VScoringMesh.hh"
#include "cnpy.h"

/**
 * @brief const G4int lin = x * Ny * Nz + y * Nz + z;  // 线性索引
 * x, y, z; z 变化最快, fortran 顺序
 *
 */

void MyNpzScoreWriter::DumpQuantityToFile(const G4String& psName, const G4String& fileName, const G4String& option)
{
    if (!fScoringMesh) return;

#if 0
    ///======================================== 保留官方 CSV 输出, for debug
    G4VScoreWriter::DumpQuantityToFile(psName, fileName, option);
#endif

    /*---------- 构造 .npz 文件名 ----------*/
    std::string npzName = fileName;
    std::size_t dot = npzName.rfind('.');
    if (dot != std::string::npos) npzName.erase(dot);
    npzName += ".npz";

    /*---------- 取 map ----------*/
    auto scoreMap = fScoringMesh->GetScoreMap();
    auto it = scoreMap.find(psName);
    if (it == scoreMap.end()) return;

    using ScoreMap = std::map<G4int, G4StatDouble*>;
    ScoreMap* score = it->second->GetMap();

    G4double unitValue = fScoringMesh->GetPSUnitValue(psName);
    const G4int Nx = fNMeshSegments[0];
    const G4int Ny = fNMeshSegments[1];
    const G4int Nz = fNMeshSegments[2];
    const G4int tot = Nx * Ny * Nz;

    /*---------- 六列缓冲区 ----------*/
    std::vector<G4int> col_x(tot), col_y(tot), col_z(tot);
    std::vector<G4double> col_val(tot), col_val2(tot), col_ent(tot);

    /*---------- 逐体素填充 ----------*/
    for (G4int x = 0; x < Nx; ++x)
    {
        for (G4int y = 0; y < Ny; ++y)
        {
            for (G4int z = 0; z < Nz; ++z)
            {
                const G4int idx = this->GetIndex(x, y, z);
                const G4int lin = x * Ny * Nz + y * Nz + z;  // 线性索引

                col_x[lin] = x;
                col_y[lin] = y;
                col_z[lin] = z;

                auto vit = score->find(idx);
                if (vit != score->end() && vit->second)
                {
                    G4StatDouble* s = vit->second;
                    col_val[lin] = s->sum_wx() / unitValue;
                    col_val2[lin] = s->sum_wx2() / unitValue / unitValue;
                    col_ent[lin] = static_cast<G4double>(s->n());
                }
                else
                {
                    col_val[lin] = 0.0;
                    col_val2[lin] = 0.0;
                    col_ent[lin] = 0.0;
                }
            }
        }
    }

    ///---------- 一次性写进 .npz ----------
    // zipname 存储文件名, fname: 变量名
    cnpy::npz_save(npzName, "idx0", col_x, "w");
    cnpy::npz_save(npzName, "idx1", col_y, "a");
    cnpy::npz_save(npzName, "idx2", col_z, "a");
    cnpy::npz_save(npzName, "value", col_val, "a");
    cnpy::npz_save(npzName, "val2", col_val2, "a");
    cnpy::npz_save(npzName, "entry", col_ent, "a");
}